package cn.stylefeng.guns.modular.gridsystem.websocket;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.collection.ConcurrentHashSet;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import cn.hutool.log.Log;
import cn.hutool.log.LogFactory;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.EqualsAndHashCode;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

@ServerEndpoint("/ws/ssh/{token}")
@Component
public class WebSocketSSHServer {

    private Log LOG = LogFactory.get();

    private static ConcurrentHashMap<String, Session> sessionMap = new ConcurrentHashMap<>();

    private static Set<KV> clientConstraint = new ConcurrentHashSet<>();

    @Data
    @EqualsAndHashCode
    @AllArgsConstructor
    public static class KV {
        private String key;
        private String value;
    }

    /**
     * 发送消息
     *
     * @param session
     * @param message
     * @throws IOException
     */
    public void sendMessage(Session session, String message) throws IOException {
        if (session != null) {
            synchronized (session) {
                session.getBasicRemote().sendText(message);
            }
        }


    }

    /**
     * 发送消息
     *
     * @param session
     * @param message
     * @throws IOException
     */
    public void sendBytes(Session session, byte[] message) throws IOException {
        if (session != null) {
            synchronized (session) {
                session.getBasicRemote().sendBinary(ByteBuffer.wrap(message));
            }
        }
    }

    /**
     * 获取session
     *
     * @param token
     * @return
     */
    public Session getSession(String token) {
        final Session session = sessionMap.get(token);
        return session;
    }

    @OnOpen
    public void onOpen(Session session, @PathParam("token") String token) {
        sessionMap.put(token, session);
        LOG.info("{}加入 websocket ssh", token);

    }

    @OnClose
    public void onClose(@PathParam("token") String token) {
        sessionMap.remove(token);
        LOG.info("{}退出 websocket ssh", token);
        final List<KV> collect = clientConstraint.parallelStream().filter(item -> StrUtil.equals(token, item.key) || StrUtil.equals(token, item.value)).collect(Collectors.toList());
        if (!CollectionUtil.isEmpty(collect)) {
            for (KV kv : collect) {
                Message message = new Message();
                if (StrUtil.equals(kv.getKey(), token)) {
                    //from断开,则ssh端要断开
                    message.setFrom(kv.getKey());
                    message.setTo(kv.getValue());
                    message.setCmd("CLOSE");
                } else {
                    //ssh端断开,发消息给client,提示断开
                    message.setFrom(kv.getValue());
                    message.setTo(kv.getKey());
                    message.setCmd("CLOSE");
                }
                final Session destination = getSession(message.getTo());
                try {
                    if (destination == null) {
                        Message empty = new Message();
                        empty.setCmd("OFFLINE");
                        empty.setFrom(message.getFrom());
                        empty.setTo(message.getTo());
                        sendMessage(getSession(message.getFrom()), JSONUtil.toJsonStr(empty));
                        return;
                    }
                    sendMessage(destination, JSONUtil.toJsonStr(message));
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }

    @OnMessage
    public void onMessage(byte[] data, Session session) {
        System.out.println("onMessage");
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        try {
            if (StrUtil.isEmpty(message.trim())) {
                //心跳包
                session.getBasicRemote().sendText("");
                return;
            }
            final Message entity = JSONUtil.toBean(message, Message.class);
            if (entity == null) {
                LOG.info("解析失败");
                return;
            }
            final String cmd = entity.getCmd();
            final Session destination = getSession(entity.getTo());
            if (destination == null) {
                Message empty = new Message();
                empty.setCmd("OFFLINE");
                empty.setFrom(entity.getTo());
                empty.setTo(entity.getFrom());
                sendMessage(session, JSONUtil.toJsonStr(empty));
                return;
            }

            if (StrUtil.equals("AUTH", cmd)) {
                //授权请求
                final String from = entity.getFrom();
                final String to = entity.getTo();
                clientConstraint.add(new KV(from, to));
            }

            sendMessage(destination, message);
//            switch (cmd) {
//                case "AUTH":
//                    //web端发起登录到client
//                    sendMessage(session, message);
//                    break;
//                case "AUTH_REPLY":
//                    //client反馈登录结果
//
//                    break;
//                case "SSH":
//                    break;
//                case "SSH_REPLY":
//                    break;
//
//            }


        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
